package com.zwh.shardingjdbc.config.readwrite.algorithm.loadbalance;

import lombok.Getter;
import org.apache.shardingsphere.infra.context.transaction.TransactionConnectionContext;
import org.apache.shardingsphere.readwritesplitting.spi.ReadQueryLoadBalanceAlgorithm;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.interceptor.TransactionAspectSupport;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.util.List;
import java.util.Properties;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 自定义基于轮询的读库负载均衡算法<br/>
 * 在类型：ROUND_ROBIN， RoundRobinReadQueryLoadBalanceAlgorithm类基础上增加自定义：<br/>
 * 当事务是只读事务 @Transactional(readOnly = true)， 查询走从库，<br/>
 * 写事务，查询走主库 <br/>
 * @author zwh
 * @date 2023-5-25
 */
public class CustomFixedReplicaRoundRobinReadQueryLoadBalanceAlgorithm implements ReadQueryLoadBalanceAlgorithm {
    private final AtomicInteger count = new AtomicInteger(0);

    @Getter
    private Properties props;

    @Override
    public void init(final Properties props) {
        this.props = props;
    }

    /*
    注意：TransactionConnectionContext上下文，只能在一个事务过程中保持。
        或者说一个http请求中，一个事务内的sql操作，读写分离会被路由到固定数据源副本中。
事务管理器
调用
AbstractPlatformTransactionManager.commit(TransactionStatus)
	AbstractPlatformTransactionManager.cleanupAfterCompletion()
		cleanupAfterCompletion(status)
			doCleanupAfterCompletion(status.getTransaction())
				DataSourceUtils.releaseConnection(con, this.dataSource)
					。。。
					con.close() // 会关闭连接，
	所以connectionContext也会被销毁, connectionContext.getTransactionConnectionContext()
     */
    @Override
    public String getDataSource(final String name, final String writeDataSourceName, final List<String> readDataSourceNames, final TransactionConnectionContext context) {
        if (context.isInTransaction()) {
            if (TransactionSynchronizationManager.isActualTransactionActive()
                    && !TransactionSynchronizationManager.isCurrentTransactionReadOnly()) {
                // 事务中，且事务@Transactional(readOnly = false) 才走主库
                return writeDataSourceName;
            } else {
                // 只读事务中，固定走同一个副本
                if (null == context.getReadWriteSplitReplicaRoute()) {
                    context.setReadWriteSplitReplicaRoute(readDataSourceNames.get(Math.abs(count.getAndIncrement()) % readDataSourceNames.size()));
                }
                return context.getReadWriteSplitReplicaRoute();
            }
        }
        return readDataSourceNames.get(Math.abs(count.getAndIncrement()) % readDataSourceNames.size());

    }

    @Override
    public String getType() {
        return "CUSTOM_FIXED_REPLICA_ROUND_ROBIN";
    }

    @Override
    public boolean isDefault() {
        return true;
    }
}
